#include "tgaimage.h"
#include <iostream>
#include "model.h"
#include <algorithm>
#include "glm.h"
#include "glProgram.h"
#include "IShader.h"

const TGAColor white = TGAColor(255, 255, 255, 255);
const TGAColor red = TGAColor(255, 0, 0, 255);
const TGAColor green = TGAColor(0, 255, 0, 255);
const TGAColor blue = TGAColor(0, 0, 255, 255);

const int width = 800;
const int height = 800;
Vec3f light_dir = Vec3f(1, -1, -1);// 光线方向
Vec3f light_color = Vec3f(1, 1, 1);// 光线的颜色
Vec3f cameraPos = Vec3f(0, 0, 2);
Vec3f targetPos = Vec3f(0, 0, 0);

float* depthMapBuffer;

struct PhongVertexInfo : VertexInfoBase
{
	Vec2f textureCoord;
	Vec3f normal;
	Vec3f tangent;

	virtual ~PhongVertexInfo() = default;
};

class PhongVertexShader : public VertexShader
{
public:
	// 确定有哪些 Varying 变量
	virtual void ConfirmVaryingVar() override
	{
		AddVaryingVar<Vec2f>("textureCoord");
		AddVaryingVar<Vec3f>("N");
		AddVaryingVar<Vec3f>("T");
		AddVaryingVar<Vec3f>("worldPos");
        AddVaryingVar<Vec4f>("posForLightSpace");
	}
	virtual Vec4f vertex(std::shared_ptr<VertexInfoBase> vertexInfos) override
	{
		// 解码顶点属性实际类型
		auto* info = static_cast<PhongVertexInfo*>(vertexInfos.get());
		// 获得矩阵
		glm::mat4 modelMatrix = GetUniformVar<glm::mat4>("Model");
		glm::mat4 viewMatrix = GetUniformVar<glm::mat4>("View");
        glm::mat4 proMatrix = GetUniformVar<glm::mat4>("Proj");
        glm::mat4 lightSpaceMatrix = GetUniformVar<glm::mat4>("lightSpaceMatrix");
		// 设置 Varying 变量
		SetVaryingVar<Vec2f>("textureCoord", info->index, info->textureCoord);
		glm::mat3 normalMatrix = modelMatrix.inverse().transpose().to_mat3();
		Vec3f N = (normalMatrix * info->normal).normalized();
		Vec3f T = (normalMatrix * info->tangent).normalized();
		Vec3f worldPos = modelMatrix * Vec4f(info->location, 1.f);
		Vec4f posForLightSpace = lightSpaceMatrix * Vec4f(worldPos, 1.f);
		
        SetVaryingVar<Vec3f>("N", info->index, N);
        SetVaryingVar<Vec3f>("T", info->index, T);
		SetVaryingVar<Vec3f>("worldPos", info->index, worldPos);
        SetVaryingVar<Vec4f>("posForLightSpace", info->index, posForLightSpace);

		return proMatrix * viewMatrix * modelMatrix * Vec4f(info->location, 1.f);
	}
};

class PhongFragmentShader : public FragmentShader
{
public:
	virtual bool fragment(FragmentInfo& info) override
	{
		TGAImage* diffuseTexture = GetUniformVar<TGAImage*>("diffuseTexture");// 纹理贴图
		TGAImage* normalTexture = GetUniformVar<TGAImage*>("normalTexture");// 法线贴图
		TGAImage* specularTexture = GetUniformVar<TGAImage*>("specularTexture");
		TGAImage* glowTexture = GetUniformVar<TGAImage*>("glowTexture");
		TGAImage* depthMap = GetUniformVar<TGAImage*>("depthMap");
		glm::mat4 modelMatrix = GetUniformVar<glm::mat4>("Model");

		Vec3f lightDir = GetUniformVar<Vec3f>("lightDir");// 光线方向
		lightDir = modelMatrix * Vec4f(lightDir, 0.f);
		lightDir.normalize();
		Vec3f lightColor = GetUniformVar<Vec3f>("lightColor");
		Vec3f cameraPos = GetUniformVar<Vec3f>("cameraPos");

		Vec2f textureCoord = GetInterVaryingVar<Vec2f>("textureCoord", info.bc_coord);// 获得插值后的纹理坐标
		Vec3f worldPos = GetInterVaryingVar<Vec3f>("worldPos", info.bc_coord);
		Vec4f posForLightSpace = GetInterVaryingVar<Vec4f>("posForLightSpace", info.bc_coord);

		Vec3f N = GetInterVaryingVar<Vec3f>("N", info.bc_coord).normalized();
		Vec3f T = GetInterVaryingVar<Vec3f>("T", info.bc_coord).normalized();
		T = (T - N * T.dot(N)).normalized();//施密特正交化
		Vec3f B = N.cross(T).normalized();

		glm::mat3 TBN = glm::Vec3toMat3(T, B, N);
		Vec3f normal = normalTexture->texture(textureCoord.x, textureCoord.y);
		normal = TBN * normal;
		normal.normalize();

		Vec3f R = glm::reflect(lightDir, N).normalized();
		Vec3f V = (cameraPos - worldPos).normalized();
		float NdotL = std::max(0.f, -normal.dot(lightDir));
		float VdotR = std::max(0.f, V.dot(R));
		TGAColor color = diffuseTexture->texture(textureCoord.x, textureCoord.y);
		Vec3f ambient{ 0.3,0.3,0.3 };
		Vec3f k_d{ 1,1,1 };
		Vec3f diffuse = k_d.Multi(lightColor) * NdotL;
		Vec3f k_s{ 0.5,0.5,0.5 };
		float shininess = specularTexture->texture(textureCoord.x, textureCoord.y).b;
		Vec3f specular = k_s.Multi(lightColor) * std::pow(VdotR, shininess);

		// 计算阴影
		float shadow = 0.f;
		{
			Vec4f projCoords = posForLightSpace / posForLightSpace.w;
			//projCoords = projCoords * 0.5 + Vec3f(0.5, 0.5, 0.5);
			//float closestDepth = depthMap->texture(projCoords.x, projCoords.y).r / 255.f;
			projCoords = m_program->GetViewPort() * projCoords;
			float closestDepth = depthMapBuffer[int(projCoords.x) + int(projCoords.y) * width];
			float currentDepth = projCoords.z;
			float bias = std::max(0.015 * (1.0 - std::max(normal.dot(lightDir), 0.f)), 0.001);
			shadow = (currentDepth > closestDepth + bias) ? 1.f : 0.f;
		}

		Vec3f I = ambient + (1 - shadow) * (diffuse + specular);

		info.color = color.Multi(I.x, I.y, I.z) + glowTexture->texture(textureCoord.x, textureCoord.y);

		return false;// do not discard the pixel
	}
};

class DepthVertexShader : public VertexShader
{
public:
	virtual void ConfirmVaryingVar() override {}
	virtual Vec4f vertex(std::shared_ptr<VertexInfoBase> vertexInfos) override
	{
		glm::mat4 modelMatrix = GetUniformVar<glm::mat4>("Model");
		glm::mat4 lightSpaceMatrix = GetUniformVar<glm::mat4>("lightSpaceMatrix");

		return lightSpaceMatrix * modelMatrix * Vec4f(vertexInfos->location, 1.f);
	}
};

class DepthFragmentShader : public FragmentShader
{
public:
	virtual bool fragment(FragmentInfo& info) override
	{
		return false;
	}
};

int main(int argc, char** argv) 
{
	// head的模型矩阵
	glm::mat4 modelMatrix;
	//modelMatrix = glm::translate(modelMatrix, Vec3f(-0.5, 0, 0));
	//modelMatrix = glm::rotate(modelMatrix, 22.5, Vec3f(0, 1, 0));
	//modelMatrix = glm::scale(modelMatrix, Vec3f(0.5, 0.5, 0.5));

	std::shared_ptr<Model> headModel = std::make_shared<Model>("obj/diablo3_pose.obj");
	std::shared_ptr<TGAImage> diffuseTexture = std::make_shared<TGAImage>(width, height, TGAImage::RGB);
	std::shared_ptr<TGAImage> normalTexture = std::make_shared<TGAImage>(width, height, TGAImage::RGB);
	std::shared_ptr<TGAImage> specularTexture = std::make_shared<TGAImage>(width, height, TGAImage::RGB);
    std::shared_ptr<TGAImage> glowTexture = std::make_shared<TGAImage>(width, height, TGAImage::RGB);
	std::shared_ptr<TGAImage> depthMap = std::make_shared<TGAImage>(width, height, TGAImage::RGB);
	diffuseTexture->read_tga_file("resource/diablo3_pose_diffuse.tga");
    normalTexture->read_tga_file("resource/diablo3_pose_nm_tangent.tga");
    specularTexture->read_tga_file("resource/diablo3_pose_spec.tga");
    glowTexture->read_tga_file("resource/diablo3_pose_glow.tga");
	diffuseTexture->flip_vertically();
    normalTexture->flip_vertically();
    specularTexture->flip_vertically();
    glowTexture->flip_vertically();
	auto* zbuffer = new float[width * height];
	depthMapBuffer = new float[width * height];
	std::fill(zbuffer, zbuffer + width * height, std::numeric_limits<float>::max());
	std::fill(depthMapBuffer, depthMapBuffer + width * height, std::numeric_limits<float>::max());
	TGAImage result(width, height, TGAImage::RGB);
	std::array<std::shared_ptr<VertexInfoBase>, 3> vertexInfos;

	// 第一遍渲染——生成 Depth Map
	std::shared_ptr<DepthFragmentShader> depthFragmentShader = std::make_shared<DepthFragmentShader>();
	std::shared_ptr<DepthVertexShader> depthVertexShader = std::make_shared<DepthVertexShader>();
	std::shared_ptr<glProgram> depthProgram =std::make_shared<glProgram>();
	depthProgram->RegisterVertexShader(depthVertexShader);
    depthProgram->RegisterFragmentShader(depthFragmentShader);
    depthProgram->SetViewPort(width, height);

	depthProgram->AddUniformVar<glm::mat4>("lightSpaceMatrix");
	depthProgram->AddUniformVar<glm::mat4>("Model");

	glm::mat4 lightProjection = glm::ortho(-5.f, 5.f, -5.f, 5.f, 1.0f, 20.f);
	glm::mat4 lightView = glm::lookAt(-light_dir * 10, Vec3f(0, 0, 0), Vec3f(0, 1, 0));
    depthProgram->SetUniformVar<glm::mat4>("lightSpaceMatrix", lightProjection * lightView);
    depthProgram->SetUniformVar<glm::mat4>("Model", modelMatrix);

	for (int i = 0; i < headModel->nfaces(); i++)
	{
		std::vector<int> face_v = headModel->face_v(i);

		std::array<std::pair<Vec3f, Vec2f>, 3> calTanInfos;
		for (int j = 0; j < 3; j++)
		{
			std::shared_ptr<VertexInfoBase> info = std::make_shared<VertexInfoBase>();
			info->location = Vec4f(headModel->vert(face_v[j]), 1.f);
			vertexInfos[j] = info;
		}

		depthProgram->Draw(result, depthMapBuffer, vertexInfos);
	}
	// 将 zbuffer 的值写入 depth map 中
    for (int i = 0; i < width * height; i++)
    {
		int x = i % width;
		int y = i / width;
        if (depthMapBuffer[i] >= 1)
        {	
			depthMap->set(x, y, white);
        }
		else
		{
			unsigned int color = depthMapBuffer[i] * 255;
            depthMap->set(x, y, TGAColor(color, color, color, 255));
		}
    }
	depthMap->flip_vertically();
    depthMap->write_tga_file("depthMap.tga");

	//-----------------------------------------------------------------------------------

	// 第二遍渲染——正常渲染
	std::shared_ptr<PhongVertexShader> phongVertexShader = std::make_shared<PhongVertexShader>();
	std::shared_ptr<PhongFragmentShader> phongFragmentShader = std::make_shared<PhongFragmentShader>();
	std::shared_ptr<glProgram> phongProgram = std::make_shared<glProgram>();

	phongProgram->RegisterVertexShader(phongVertexShader);
    phongProgram->RegisterFragmentShader(phongFragmentShader);
	phongProgram->SetViewPort(width, height);

	phongProgram->AddUniformVar<glm::mat4>("Model");
    phongProgram->AddUniformVar<glm::mat4>("View");
    phongProgram->AddUniformVar<glm::mat4>("Proj");
    phongProgram->AddUniformVar<Vec3f>("lightDir");
	phongProgram->AddUniformVar<Vec3f>("lightColor");
	phongProgram->AddUniformVar<TGAImage*>("diffuseTexture");
    phongProgram->AddUniformVar<TGAImage*>("normalTexture");
    phongProgram->AddUniformVar<TGAImage*>("specularTexture");
    phongProgram->AddUniformVar<TGAImage*>("glowTexture");
	phongProgram->AddUniformVar<Vec3f>("cameraPos");
	phongProgram->AddUniformVar<TGAImage*>("depthMap");
	phongProgram->AddUniformVar<glm::mat4>("lightSpaceMatrix");

	glm::mat4 viewMatrix = glm::lookAt(cameraPos, targetPos, Vec3f(0, 1, 0));
	glm::mat4 projMatrix = glm::perspective(45, float(width) / height, 0.1f, 100.f);
    phongProgram->SetUniformVar<glm::mat4>("Model", modelMatrix);
    phongProgram->SetUniformVar<glm::mat4>("View", viewMatrix);
    phongProgram->SetUniformVar<glm::mat4>("Proj", projMatrix);
    phongProgram->SetUniformVar<Vec3f>("lightDir", light_dir);
	phongProgram->SetUniformVar<Vec3f>("lightColor", light_color);
    phongProgram->SetUniformVar<TGAImage*>("diffuseTexture", diffuseTexture.get());
    phongProgram->SetUniformVar<TGAImage*>("normalTexture", normalTexture.get());
    phongProgram->SetUniformVar<TGAImage*>("specularTexture", specularTexture.get());
    phongProgram->SetUniformVar<TGAImage*>("glowTexture", glowTexture.get());
    phongProgram->SetUniformVar<Vec3f>("cameraPos", cameraPos);
    phongProgram->SetUniformVar<TGAImage*>("depthMap", depthMap.get());
	phongProgram->SetUniformVar<glm::mat4>("lightSpaceMatrix", lightProjection * lightView);

	for (int i = 0; i < headModel->nfaces(); i++)
	{
		std::vector<int> face_v = headModel->face_v(i);
		std::vector<int> face_vt = headModel->face_vt(i);
		std::vector<int> face_vn = headModel->face_vn(i);

		std::array<std::pair<Vec3f, Vec2f>, 3> calTanInfos;
		for (int j = 0; j < 3; j++)
		{
			std::shared_ptr<PhongVertexInfo> info = std::make_shared<PhongVertexInfo>();
			info->location = Vec4f(headModel->vert(face_v[j]), 1.f);
			info->textureCoord = headModel->vertTexture(face_vt[j]);
			info->normal = headModel->vertNormal(face_vn[j]);

			vertexInfos[j] = info;
			calTanInfos[j] = std::make_pair(info->location, info->textureCoord);
		}
		// 计算切线
		Vec3f t = glm::tangent(calTanInfos[0], calTanInfos[1], calTanInfos[2]).normalized();
		for (int j = 0; j < 3; j++)
		{
			auto* info = static_cast<PhongVertexInfo*>(vertexInfos[j].get());
			info->tangent = t;
		}

		phongProgram->Draw(result, zbuffer, vertexInfos);
	}

	for (int i = 0;i < 3;i++)
	{
		vertexInfos[i] = nullptr;
	}
	result.flip_vertically();
	result.write_tga_file("output.tga");

	delete[] zbuffer;
    delete[] depthMapBuffer;
	return 0;
}
